from data_reader import *
from neural_net import *
from neural_net_impl_sol import *
import sys
import random

class SimpleNetwork:
  def __init__(self):
    self.network = NeuralNetwork()
    for i in range(14):
      for j in range(14):
        node = Node()
        self.network.AddNode(node, NeuralNetwork.INPUT)
    for i in range(10):
      node = Node()
      for input in self.network.inputs:
        node.AddInput(input, None, self.network)
      self.network.AddNode(node, NeuralNetwork.OUTPUT)

  def EncodeLabel(self, label):
    """Returns the target for the given label."""
    target = Target()
    for i in range(10):
      if label == i:
        target.values.append(1.0)
      else:
        target.values.append(0.0)
    return target

  def GetNetworkLabel(self):
    """Returns the label based on the network's current output values."""
    best_score = -1
    best_label = -1
    for (output, label) in zip(self.network.outputs, range(len(self.network.outputs))):
      if best_label == -1 or output.transformed_value > best_score:
        best_score = output.transformed_value
        best_label = label
    return best_label

  def Convert(self, image):
    input = Input()
    for row in image.pixels:
      for entry in row:
        input.values.append(entry / 256.0)
    return input

  def Classify(self, image):
    input = self.Convert(image)
    FeedForward(self.network, input)
    return self.GetNetworkLabel()

  def Performance(self, images):
    correct = 0
    for image in images:
      if self.Classify(image) == image.label:
        correct += 1
    return correct * 1.0 / len(images)

  def Train(self, images, validation_images, learning_rate, epochs):
    inputs = []
    targets = []
    for image in images:
      inputs.append(self.Convert(image))
      targets.append(self.EncodeLabel(image.label))
    for i in range(epochs):
      Train(self.network, inputs, targets, learning_rate, 1)
      print '%d Performance: %.8f %.3f' % (
        i + 1, self.Performance(images), self.Performance(validation_images))

def FillArgsMap(argv, args_map):
  curkey = None
  for i in range(1, len(args_map)):
    key = args_map[i]
    if key.find('-') == 0:
      args_map[key] = True
      curkey = key
    else:
      assert curkey
      args_map[curkey] = key

def main():
  args_map = {}
  args_map['-e'] = 10
  args_map['-r'] = 1.0
  network = SimpleNetwork()
  epochs = int(args_map['-e'])
  rate = float(args_map['-r'])
  print 'Epochs: %d, Learning Rate: %f' % (epochs, rate)
  images = DataReader.GetImages('training-9k.txt', -1)
  for image in images:
    assert len(image.pixels) == 14
    assert len(image.pixels[0]) == 14

  validation = DataReader.GetImages('validation-1k.txt', -1)
  # Initialize weights to small random weights
  for weight in network.network.weights:
    weight.value = random.random() * 0.00001
  network.Train(images, validation, rate, epochs)

main()
